/*
 * CRF.cpp
 *
 *  Created on: Sep 7, 2012
 *      Author: v-suqil
 */

#include "CRF.h"

CRF::CRF(InstanceList* trainlist) {
	this->modellist = trainlist;
	this->numvars = 0;
	this->numStates = this->modellist->getLabelAlphabet()->size();
	this->pars = new par_t*[numStates];
	this->inipars = new par_t[numStates];
	this->finpars = NULL;
	for (int i = 0; i < numStates; ++i) {
		this->pars[i] = new par_t[numStates];
	}
	this->vars = NULL;
	this->grads = NULL;
	build();
}

void CRF::build() {
	Instance* inst;
	value_t* labels;
	quiver_t* fvs;
	int length;
	for (unsigned i = 0; i < this->modellist->size(); ++i) {
		inst = this->modellist->get(i);
		labels = inst->getLabels();
		fvs = inst->getFeatureVectors();
		length = inst->size();

		for (unsigned k = 0; k < fvs->at(0)->size(); ++k) {
			++this->inipars[labels->at(0)][fvs->at(0)->at(k)].weight;
		}
		for (int j = 1; j < length; ++j) {
			for (unsigned k = 0; k < fvs->at(j)->size(); ++k) {
				++this->pars[labels->at(j)][labels->at(j - 1)][fvs->at(j)->at(k)].weight;
			}
		}
		if (finpars != NULL) {
			for (unsigned k = 0; k < fvs->at(length - 1)->size(); ++k) {
				++this->finpars[labels->at(length - 1)][fvs->at(length - 1)->at(
						k)].weight;
			}
		}
	}

//	for (int i = 0; i < numStates; ++i) {
//		for (int j = 0; j < numStates; ++j) {
//			for (par_t::iterator it = pars[i][j].begin();
//					it != pars[i][j].end(); ++it) {
//				printf("%d ", (*it).second.weight);
//			}
//		}
//	}
}

void CRF::train() {
	train(10000);
}

void CRF::train(int maxiter) {
	printf("Start training...\n");
	SCRF::LBFGS lbfgs;
	double of;
	double f = 0;
	unsigned int converge = 0;
	for (int i = 0; i < maxiter; ++i) {
		printf("No.%d iteration... ", i);
		of = f;
		f = getValue();
		getVariables();
		getGradients();
//		f -= gaussian();
		printf("Value = %f\n", f);
		if (2 * std::abs(f - of) < 0.00001 * (std::abs(f) + std::abs(of)))
			++converge;
		if (converge == 3) break;
		lbfgs.optimize(getNumVars(), vars->data(), f, grads->data(), false,
				1.0);
		setParameters();
	}
}

int CRF::getNumStates() {
	return this->numStates;
}

var_t* CRF::getVariables() {
	if (this->vars == NULL) {
		this->vars = new var_t();
	} else {
		this->vars->clear();
	}
	for (int i = 0; i < numStates - 1; ++i) {
		for (par_t::iterator it = inipars[i].begin(); it != inipars[i].end();
				++it) {
			vars->push_back((*it).second.value);
		}
	}
	for (int i = 0; i < numStates; ++i) {
		for (int j = 0; j < numStates; ++j) {
			for (par_t::iterator it = pars[i][j].begin();
					it != pars[i][j].end(); ++it) {
				vars->push_back((*it).second.value);
			}
		}
	}
	if (finpars != NULL) {
		for (int i = 0; i < numStates; ++i) {
			for (par_t::iterator it = finpars[i].begin();
					it != finpars[i].end(); ++it) {
				vars->push_back((*it).second.value);
			}
		}
	}
	return vars;
}

void CRF::setParameters() {
	int c = 0;
	for (int i = 0; i < numStates; ++i) {
		for (par_t::iterator it = inipars[i].begin(); it != inipars[i].end();
				++it) {
			(*it).second.value = vars->at(c);
			++c;
		}
	}
	for (int i = 0; i < numStates; ++i) {
		for (int j = 0; j < numStates; ++j) {
			for (par_t::iterator it = pars[i][j].begin();
					it != pars[i][j].end(); ++it) {
				(*it).second.value = vars->at(c);
				++c;
			}
		}
	}
	if (finpars != NULL) {
		for (int i = 0; i < numStates; ++i) {
			for (par_t::iterator it = finpars[i].begin();
					it != finpars[i].end(); ++it) {
				(*it).second.value = vars->at(c);
				++c;
			}
		}
	}
}

var_t* CRF::getGradients() {
	if (this->grads == NULL) {
		this->grads = new var_t();
	} else {
		this->grads->clear();
	}
	feat* p;
	for (int i = 0; i < numStates; ++i) {
		for (par_t::iterator it = inipars[i].begin(); it != inipars[i].end();
				++it) {
			p = &(*it).second;
			grads->push_back(p->weight - p->expect - p->value / GV);
		}
	}
	for (int i = 0; i < numStates; ++i) {
		for (int j = 0; j < numStates; ++j) {
			for (par_t::iterator it = pars[i][j].begin();
					it != pars[i][j].end(); ++it) {
				p = &(*it).second;
				grads->push_back(p->weight - p->expect - p->value / GV);
			}
		}
	}
	if (finpars != NULL) {
		for (int i = 0; i < numStates; ++i) {
			for (par_t::iterator it = finpars[i].begin();
					it != finpars[i].end(); ++it) {
				p = &(*it).second;
				grads->push_back(p->weight - p->expect - p->value / GV);
			}
		}
	}
	return grads;
}

double CRF::totalValue(Instance* inst) {
	int length = inst->size();
	quiver_t* input = inst->getFeatureVectors();
	double** alpha = new double*[length];
	double** beta = new double*[length];
	double*** weights = new double**[length];
	for (int ip = 0; ip < length; ++ip) {
		if (ip != 0) {
			weights[ip] = new double*[numStates];
			for (int i = 0; i < numStates; ++i) {
				weights[ip][i] = new double[numStates];
				for (int j = 0; j < numStates; ++j) {
					weights[ip][i][j] = getWeight(input->at(ip), i, j);
				}
			}
		}
		alpha[ip] = new double[numStates];
		beta[ip] = new double[numStates];
	}
	//initialize forward variables
	for (int i = 0; i < numStates; ++i) {
		alpha[0][i] = getInitialWeight(input->at(0), i);
	}
	//forward
	for (int ip = 1; ip < length; ++ip) {
		for (int i = 0; i < numStates; ++i) {
			alpha[ip][i] = NEG_INF;
			for (int j = 0; j < numStates; ++j) {
				alpha[ip][i] = sumLog(alpha[ip][i],
						weights[ip][i][j] + alpha[ip - 1][j]);
			}
		}
	}

	double totalWeight = NEG_INF;
	for (int i = 0; i < numStates; ++i) {
		totalWeight = sumLog(totalWeight, alpha[length - 1][i]);
	}
	//backward
	par_t::iterator it;
	//initialize backward variables
	for (int i = 0; i < numStates; ++i) {
		beta[length - 1][i] = getFinalWeight(input->at(length - 1), i);
		double prob = exp(
				alpha[length - 1][i] + beta[length - 1][i] - totalWeight);
		if (finpars != NULL) {
			for (unsigned k = 0; k < input->at(length - 1)->size(); ++k) {
				it = finpars[i].find(input->at(length - 1)->at(k));
				if (it != finpars[i].end()) {
					(*it).second.expect += prob;
				}
			}
		}
	}
	for (int ip = length - 2; ip > 0; --ip) {
		for (int j = 0; j < numStates; ++j) {
			beta[ip][j] = NEG_INF;
			for (int i = 0; i < numStates; ++i) {
				beta[ip][j] = sumLog(beta[ip][j],
						weights[ip + 1][i][j] + beta[ip + 1][i]);
			}
			for (int i = 0; i < numStates; ++i) {
				double prob = exp(
						alpha[ip - 1][i] + weights[ip][i][j] + beta[ip][j]
								- totalWeight);
				for (unsigned k = 0; k < input->at(ip)->size(); ++k) {
					it = pars[i][j].find(input->at(ip)->at(k));
					if (it != pars[i][j].end()) {
						(*it).second.expect += prob;
					}
				}
			}
		}
	}
	//free memory
	for (int ip = 0; ip < length; ++ip) {
		if (ip != 0) {
			for (int i = 0; i < numStates; ++i) {
				delete weights[ip][i];
			}
			delete weights[ip];
		}
		delete alpha[ip];
		delete beta[ip];
	}
	delete alpha;
	delete beta;
	delete weights;
	return totalWeight;
}

double CRF::constrainValue(Instance* inst) {
	double constrain(0.0);
	value_t* lbs = inst->getLabels();
	quiver_t* fvs = inst->getFeatureVectors();
	int length = inst->size();
	for (unsigned k = 0; k < fvs->at(0)->size(); ++k) {
		constrain += this->inipars[lbs->at(0)][fvs->at(0)->at(k)].value;
	}
	for (int j = 1; j < length; ++j) {
		for (unsigned k = 0; k < fvs->at(j)->size(); ++k) {
			constrain += this->pars[lbs->at(j)][lbs->at(j - 1)][fvs->at(j)->at(
					k)].value;
		}
	}
	if (finpars != NULL) {
		for (unsigned k = 0; k < fvs->at(length - 1)->size(); ++k) {
			constrain +=
					this->finpars[lbs->at(length - 1)][fvs->at(length - 1)->at(k)].value;
		}
	}
	return constrain;
}

double CRF::getWeight(value_t* featvec, int i, int j) {
	double weight = 0;
	par_t::iterator it;
	for (unsigned k = 0; k < featvec->size(); ++k) {
		it = pars[j][i].find(featvec->at(k));
		if (it == pars[j][i].end())
			continue;
		weight += (*it).second.value;
	}
	return weight;
}

double CRF::getInitialWeight(value_t* featvec, int i) {
	double weight = 0;
	par_t::iterator it;
	for (unsigned k = 0; k < featvec->size(); ++k) {
		it = inipars->find(featvec->at(k));
		if (it == inipars->end())
			continue;
		weight += inipars->at(featvec->at(k)).value;
	}
	return weight;
}

double CRF::getFinalWeight(value_t* featvec, int i) {
	double weight = 0;
	if (finpars == NULL)
		return 0;
	par_t::iterator it;
	for (unsigned k = 0; k < featvec->size(); ++k) {
		it = finpars->find(featvec->at(k));
		if (it == finpars->end())
			continue;
		weight += (*it).second.value;
	}
	return weight;
}

double CRF::getValue() {
	Instance* inst;
	double value(0.0);
	for (unsigned i = 0; i < modellist->size(); ++i) {
		inst = modellist->get(i);
		value += constrainValue(inst) - totalValue(inst);
//		value += totalValue(inst) - constrainValue(inst);
	}
	return value;
}

double CRF::gaussian() {
	double reg(0.0);
	for (unsigned i = 0; i < vars->size(); ++i) {
		reg += vars->at(i) * vars->at(i) / (2 * GV);
	}
	return reg;
}

double CRF::l1() {
	double reg(0.0);
	for (unsigned i = 0; i < vars->size(); ++i) {
		reg += lone * std::abs(vars->at(i));
	}
	return reg;
}

par_t** CRF::getParameters() {
	return this->pars;
}

par_t* CRF::getInitialParameters() {
	return this->inipars;
}

par_t* CRF::getFinalParameters() {
	return this->finpars;
}

unsigned CRF::getNumVars() {
	return vars->size();
}

CRF::~CRF() {
	for (int i = 0; i < numStates; ++i) {
		delete pars[i];
	}
	delete pars;
	delete inipars;
	delete finpars;
}

